﻿using System;
using System.Collections.Generic;

namespace PortableTPL
{
    class StandardTaskScheduler : TaskScheduler
    {
        readonly WorkerThread[] workerThreads;

        internal StandardTaskScheduler()
        {
            workerThreads = new WorkerThread[ProcessorCount];

            for (int index = 0; index < ProcessorCount; index++)
            {
                workerThreads[index] = new WorkerThread(index, this);
            }
        }

        public override void Dispose()
        {
            for (int index = 0; index < workerThreads.Length; index++)
            {
                workerThreads[index].Dispose();
            }
        }

        public override void EnqueueTask(Task task)
        {
            WorkerThread workerThread = FindLowestQueue();
            workerThread.AddTask(task);
        }

        WorkerThread FindLowestQueue()
        {
            WorkerThread result = null;
            int currentLevel = int.MaxValue;

            foreach (WorkerThread workerThread in workerThreads)
            {
                if (workerThread.QueueLength < currentLevel)
                {
                    currentLevel = workerThread.QueueLength;
                    result = workerThread;
                }
            }

            return result;
        }

        internal bool StealWork(WorkerThread source)
        {
            foreach (WorkerThread workerThread in workerThreads)
            {
                if (workerThread == source)
                    continue;

                if (workerThread.QueueLength > 1)
                {
                    int total = workerThread.QueueLength / 2;
                    List<Task> stolenWork = new List<Task>();
                    for (int index = 0; index < total; index++)
                    {
                        Task task;
                        if (workerThread.TryDequeue(out task))
                        {
                            stolenWork.Add(task);
                        }
                    }

                    source.AddTasks(stolenWork.ToArray());
                    return true;
                }
            }

            return false;
        }

        internal void Dispatch(int inclusiveStart, int exclusiveEnd, Action<int> action, ParallelOptions parallelOptions = null)
        {
            int stackSize = (exclusiveEnd - inclusiveStart) / ProcessorCount;

            if (parallelOptions != null && parallelOptions.CancellationToken != null)
                Dispatch(stackSize, exclusiveEnd, i =>
                {
                    return new Task(() => action(i), parallelOptions.CancellationToken);
                });
            else
                Dispatch(stackSize, exclusiveEnd, i =>
                {
                    return new Task(() => action(i));
                });
        }

        internal void Dispatch<T>(IList<T> list, Action<T> action, ParallelOptions parallelOptions = null)
        {
            int stackSize = (list.Count) / ProcessorCount;
            if (parallelOptions != null && parallelOptions.CancellationToken != null)
                Dispatch(stackSize, list.Count, i =>
                {
                    return new Task(() => action(list[i]), parallelOptions.CancellationToken);
                });
            else
                Dispatch(stackSize, list.Count, i =>
                {
                    return new Task(() => action(list[i]));
                });
        }

        void Dispatch(int stackSize, int excludedEnd, Func<int, Task> produceTask)
        {
            List<Task> totalTasks = new List<Task>();

            int index = 0;
            foreach (WorkerThread workerThread in workerThreads)
            {
                int start = index * stackSize;
                int end = start + stackSize;

                if (index == ProcessorCount - 1)
                    end = excludedEnd;

                Task[] tasks = new Task[end - start];
                int taskID = 0;

                for (int local = start; local < end; local++)
                {
                    Task task = produceTask(local);
                    tasks[taskID++] = task;
                }

                totalTasks.AddRange(tasks);
                workerThread.AddTasks(tasks);
                index++;
            }

            Task.WaitAll(totalTasks.ToArray());
        }
    }
}
